Replace Qwen3.5 FLA paths with Attention Gym by drisspg · Pull Request #4389 · pytorch/torchtitan

drisspg · 2026-09-01T04:45:03Z

Human Note

Lots of data; we added the paged form to conv for q_tokens > 1 and also added supprot for chunk_gdn that works on blackwell and hopper, verifying on hopper now

Agent note

This incorporates Angel Li's Attention Gym migration from #4363 and extends it so the resulting
change can land as one self-contained PR against main. It uses Attention Gym for Qwen3.5 GDN
training, batch-invariant backward recomputation, recurrent execution, decode, and paged prefill.

The vLLM path advances both convolution and SSM cache slots in place. This removes the prior
prefill-time temporary state allocations, grouped-head expansion, gather/scatter copies, transposes,
and host-syncing continuation check while preserving the shared [slots, H, V, K] SSM layout.
The paged operations landed in Attention Gym as
meta-pytorch/attention-gym#431 and
meta-pytorch/attention-gym#432.

The now-unused FLA backend selection and direct dependency are removed. Q/K normalization is explicit
at the Attention Gym boundary, and debug configurations use the fused backend's K = V = 128
contract while retaining two key heads for TP=2 coverage.

Performance

End-to-end vLLM: before either PR vs combined candidate

GB200, Qwen3.5 debug model with production GDN dimensions K=V=128, prompt length 128, generation
length 64, TP=1, FULL_AND_PIECEWISE CUDA graphs, five measured rounds per engine launch. Three engine
launches were paired at the parent of #4363 (9aadc3616) and the combined candidate; the table reports
median throughput across launches and the median paired change.

Batch	Before either PR output tok/s	Combined output tok/s	Median paired change
1	449.3	482.9	+7.5%
2	869.2	936.0	+12.3%
4	1683.9	1786.8	+6.1%
8	3184.6	3389.6	+5.3%
16	5874.8	6251.2	+6.4%

Process-level throughput varied, but every paired before/after launch improved at every batch size.

Hopper end-to-end vLLM: FLA parent vs combined candidate (September 1, 2026)

H100 80GB (SM90), Qwen3.5 debug model, BF16, prompt length 128, generation length 64,
TP=1, FULL_AND_PIECEWISE CUDA graphs, one warmup and five synchronized measured rounds per
batch. The process-level A-B-A comparison used the pre-PR FLA parent (9aadc3616, FLA 0.5.2)
and the combined candidate with Attention Gym e429589.

The parent debug config defaults to K=V=64, so the baseline was explicitly rebuilt with
K=V=128 to match the candidate and production GDN dimensions. This compares the complete FLA
GDN + convolution stack with the Attention Gym paged GDN + paged convolution stack, rather than
an isolated GDN kernel. The Attention Gym column gives the range of the two bracketing candidate
medians; change uses their mean.

Batch	FLA output tok/s	Attention Gym output tok/s	Bracket-mean change
1	431.36	435.14-445.99	+2.13%
2	840.71	861.75-866.95	+2.81%
4	1618.53	1630.57-1697.46	+2.81%
8	3014.12	2997.69-3155.29	+2.07%
16	5496.84	5419.09-5641.72	+0.61%

The matched Hopper result is roughly tied: bracket means favor Attention Gym by 0.61-2.81%, but
process-level variation overlaps FLA at batches 8 and 16. The H100 validation also covered
paged-state boundaries, mixed prefill/decode scheduling, trainer/generator bitwise parity, fused
GDN backward, an AdamW update, and post-update weight synchronization into vLLM.

Hopper training convergence (September 1, 2026)

A deterministic 100-step Qwen3.5 trainer run used the same K=V=128 architecture, exact initial
weights, fixed packed C4 batches, AdamW schedule, gradient clipping, and held-out evaluation every
10 steps. Both models had 138,074,416 parameters and the same initial-state and corpus SHA256 hashes.

Metric	FLA	Attention Gym	Attention Gym - FLA
Initial held-out loss	12.966858	12.966884	+0.000026
Final train loss	4.584238	4.584461	+0.000223
Final held-out loss	5.129257	5.129730	+0.000473
Median step time, excluding compile	168.10 ms	139.94 ms	-28.16 ms
Median tokens/s, excluding compile	1510.98	1815.01	+20.12%

Train-loss correlation was 0.99999999496 with maximum absolute delta 0.001055; held-out-loss
correlation was 0.99999999772 with maximum delta 0.000473. Gradient-norm correlation was
0.99995874, all metrics remained finite, and the final gradient-norm delta was zero. This found
no meaningful bounded convergence regression while exercising real forward, fused GDN backward,
global clipping, AdamW updates, and held-out evaluation.

This is a debug-model C4 convergence/regression run, not long-horizon production RL convergence.

GDN prefill ablation

GB200, fixed-pointer warm-cache, HK=8, H=16, K=V=128, BF16 QKV, FP32 state, five interleaved
rounds of 30 iterations:

Workload	Before: FLA + copies	Attention Gym + copies	Attention Gym direct paged	Direct vs before
N=1, T=1024	105.4 us	124.5 us	89.5 us	-15.1%
N=8, T=1024	109.2 us	110.2 us	86.9 us	-20.4%
N=8, T=4096	221.6 us	183.3 us	153.1 us	-30.9%
N=32, T=4096	314.3 us	269.0 us	203.3 us	-35.3%
N=64, T=8192	574.7 us	466.8 us	349.5 us	-39.2%

Convolution prefill ablation

At the Qwen3.5 TP=2-local shape (C=4096, W=4, BF16), direct paging reduced GPU time by
43.8-79.4% for fresh prompts and 46.2-80.4% for resumed prefixes across N=1-64 and 128-512
tokens per sequence.

Paged-state continuity

The non-batch-invariant vLLM path passed a 15-case boundary sweep at 528, 1056, and 1584 tokens
with continuation lengths 1, 63, 64, 65, and 527. Every cached/fresh pair generated identical token
IDs; the maximum logprob delta was 1.713e-3. A batch mixing fresh and resumed prefills passed, and
a forced mixed decode+prefill step matched standalone token IDs with maximum deltas 1.385e-3 and
9.537e-7.

Perfetto traces

CUDA Graph replay traces use enable_annotations=True and nested mark_kernels scopes, then
merge-traces --align to place before/after variants on labeled process tracks:

Eager execution traces:

Test Plan

# Attention Gym paged convolution, paged GDN, graph/fake-tensor, and namespace coverage
PYTHONPATH=/home/drisspg/meta/attention-gym .venv/bin/pytest -q -n 6 test/test_short_conv_cute.py test/linear/test_gdn_chunk_fused.py test/test_namespaces.py
# 151 passed, 5 skipped from source; 151 passed, 5 skipped from the built wheel

# Full Attention Gym suite excluding two independently reproduced pre-existing failures
pytest -q -n 6 --deselect=test/test_sandwich.py::test_sandwich_flex_flash_backend --deselect=test/test_selected_attention_matches_csa_reference.py::test_indexer_loss_fullgraph_forward_backward\[device0\]
# 1189 passed, 16 skipped

# Focused TorchTitan Qwen3.5 coverage
PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x
# 9 passed

# Hopper H100 production-path smoke: paged convolution, paged GDN prefill, recurrent decode
PYTHONPATH=/home/dev/meta/torchtitan-pr4389-hopper:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/smoke_pr4389_hopper_vllm.py
# Generated four tokens with finite logprobs through the production vLLM path

# Hopper H100 prefix-cache boundary sweep and mixed scheduling with CUDA graphs
PYTHONPATH=/home/dev/meta/torchtitan-pr4389-hopper:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/test_pr4389_hopper_continuity.py
PYTHONPATH=/home/dev/meta/torchtitan-pr4389-hopper:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/test_pr4389_hopper_mixed.py
# 15/15 cached/fresh boundary cases passed with identical IDs and max logprob delta 0.0
# Mixed fresh/resumed prefill matched IDs (max deltas 7.830e-4 and 2.785e-4)
# Concurrent decode+prefill matched IDs with max delta 0.0

# Hopper H100 TP=1 trainer/vLLM parity plus backward, optimizer step, and weight sync
PYTHONPATH=/home/dev/meta/torchtitan-pr4389-hopper:/home/dev/meta/attention-gym HF_ASSETS_PATH=/home/dev/meta/torchtitan-pr4389-hopper/tests/assets/tokenizer /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 -m pytest -q -s agent_space/test_pr4389_hopper_bitwise_tp1.py
# 5 passed, 4 inherited multi-GPU skips; all reported comparisons were bitwise equal
# Includes fused GDN backward, AdamW update, weight sync, cache reset, and post-update parity

# Hopper H100 matched end-to-end A-B-A throughput: candidate, FLA parent, candidate
(cd /home/dev/meta/torchtitan-pr4389-hopper && PYTHONPATH=$PWD:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/benchmark_pr4389_hopper.py)
(cd /home/dev/meta/torchtitan-pr4389-fla-hopper && PYTHONPATH=$PWD /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/benchmark_pr4389_fla_hopper.py)
(cd /home/dev/meta/torchtitan-pr4389-hopper && PYTHONPATH=$PWD:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/benchmark_pr4389_hopper.py)
# Matched K=V=128 bracket means favored Attention Gym by 0.61-2.81%; B8/B16 overlapped

# Hopper H100 deterministic 100-step C4 convergence, matched K=V=128
(cd /home/dev/meta/torchtitan-pr4389-fla-hopper && PYTHONPATH=$PWD /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 /home/dev/meta/torchtitan-pr4389-hopper/agent_space/train_convergence_pr4389.py --variant fla --repo $PWD --output-dir /home/dev/meta/torchtitan-pr4389-hopper/agent_space/convergence_k128_c4_100 --steps 100 --batch-size 2 --sequence-length 128 --train-sequences 32 --eval-sequences 8 --eval-interval 10 --learning-rate 3e-4 --warmup-steps 10 --wandb-group pr4389-h100-k128-c4-20260901)
(cd /home/dev/meta/torchtitan-pr4389-hopper && PYTHONPATH=$PWD:/home/dev/meta/attention-gym /home/dev/.venvs/torchtitan-pr4389-hopper/bin/torchrun --standalone --nproc-per-node=1 agent_space/train_convergence_pr4389.py --variant attention_gym --repo $PWD --output-dir agent_space/convergence_k128_c4_100 --steps 100 --batch-size 2 --sequence-length 128 --train-sequences 32 --eval-sequences 8 --eval-interval 10 --learning-rate 3e-4 --warmup-steps 10 --wandb-group pr4389-h100-k128-c4-20260901)
# Identical 138,074,416-parameter initial states and token corpora
# Final train loss: 4.584238 FLA vs 4.584461 Attention Gym
# Final held-out loss: 5.129257 FLA vs 5.129730 Attention Gym
# The unchanged two-role Monarch RL launcher was not run on Hopper because this host has one GPU

# Two-GPU trainer/vLLM parity with grouped-head TP=2 GDN
PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=2 -m pytest -q torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen
# 4 passed from source and again from the built wheel; all nine comparisons are bitwise identical

# Non-batch-invariant prefix-cache continuation and mixed decode+prefill
PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=1 artifacts/pr4389/test_vllm_paged_continuity.py
# 15/15 boundary cases passed; cached/fresh generated IDs matched; max delta=1.713e-3;
# mixed fresh/resumed prefill and numerically checked mixed decode+prefill passed

# Two-GPU non-batch-invariant RL integration smoke
python -m torchtitan.experiments.rl.train --module alphabet_sort --config rl_grpo_qwen3_5_debug_varlen --async-loop.num-training-steps 10 --async-loop.training-sample-builder.no-drop-zero-std-reward-groups ...
# 10/10 steps completed against both the source checkout and isolated built wheel,
# including rollout, prefill/decode, forward/backward, and weight sync

tianyu-l · 2026-09-02T04:15:16Z

+            key_head_dim=128,
+            value_head_dim=128,


is this because of minimal block size that works with attn-gym kernels? Feels Blackwell only?

yeah this is restriction for our kernels right now; while originally motivated by blackwell this will take some more work if we want to open up

wwwjn · 2026-09-02T18:31:34Z

        self.num_speculative_tokens = (
            speculative_config.num_speculative_tokens if speculative_config else 0
        )
+        if self.num_speculative_tokens != 0:


How hard to make a kernel support speculative decoding? I want to learn what features does a kernel need to support this

just now looking at this;

it looks like vllm basically adds (n_speculated_tokens, state) so expands the state cache so the potential tokens can update their slots. If we only accept M tokens then the N_speculated - M gets drop and then I think copy over M to be the base;

This is doable, would require a decent amount of changes but forsure doable feel free to open an issue in attn_gym

Use Attention Gym for Qwen3.5 GDN training, batch-invariant backward recomputation, recurrent execution, decode, and paged chunk prefill. The paged prefill path advances vLLM SSM cache slots in place, eliminating the previous gather/chunk/scatter copies and preserving the shared [slots, H, V, K] layout across prefill and decode. Remove the now-unused FLA backend selection and dependency, normalize Q/K explicitly at the Attention Gym boundary, and update debug configurations to the fused backend K=V=128 contract while retaining two key heads for TP=2 coverage. On GB200, isolated paged-prefill measurements reduced GPU time by 18.2-31.7% across five T=1024-8192, N=1-64 cases. A matched end-to-end vLLM debug-model comparison at prompt=128, generation=64, and batch sizes 1-16 improved output throughput by 1.5-6.4%. ```bash PYTHONPATH=/home/drisspg/meta/attention-gym-paged-gdn:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/attention-gym-paged-gdn:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/cpu/test_state_dict_keys.py tests/unit_tests/cpu/test_integration_test_definitions.py tests/unit_tests/cpu/test_train_spec.py -x PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym-paged-gdn torchrun --nproc-per-node=2 -m pytest torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen -v -s ```

## Human Note ## Agent note Use Attention Gym's paged causal convolution API for vLLM multi-token prefill. The operation reads fresh or resumed cache slots directly and advances them in place, removing the remaining temporary history allocation, gather, host-syncing continuation check, and final-state scatter from Qwen3.5 serving. ## Performance On GB200 at the TP=2-local Qwen3.5 shape (C=4096, W=4, BF16), direct paging reduced convolution prefill GPU time by 43.8-79.4% for fresh prompts and 46.2-80.4% for resumed prefixes across N=1-64 and 128-512 tokens per sequence. ## Test Plan ```bash PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/attention-gym .venv/bin/pytest -q -n 6 test/test_short_conv_cute.py test/linear/test_gdn_chunk_fused.py test/test_namespaces.py PYTHONPATH=/home/drisspg/meta/torchtitan:/home/drisspg/meta/attention-gym torchrun --nproc-per-node=2 -m pytest -q torchtitan/experiments/rl/tests/test_bitwise_parity.py::TestBitwiseParityQwen35DebugVarlen ```

## Human Note ## Agent note Remove the redundant subprocess import test, narrow the recurrent backward smoke check to the decay gradient it is intended to protect, and use the current expected vLLM batch-invariance import path directly instead of carrying compatibility logic for an older package layout. ## Test Plan ```bash PYTHONPATH=/home/drisspg/meta/attention-gym:/home/drisspg/meta/torchtitan pytest -q tests/unit_tests/gpu/test_qwen3_5_deltanet.py tests/unit_tests/test_qwen3_5_mrope_positions.py -x PYTHONPATH=/home/drisspg/meta/torchtitan python -c 'from torchtitan.experiments.rl.batch_invariance import patch_bmm_for_batch_invariance; patch_bmm_for_batch_invariance()' ```

pytorch-bot Bot added the ciflow/rl label Sep 1, 2026

meta-cla Bot added the CLA Signed This label is managed by the Meta Open Source bot. label Sep 1, 2026

drisspg changed the base branch from gdn_attn_gym to main September 1, 2026 04:46

drisspg force-pushed the gdn-attn-gym-paged branch 2 times, most recently from e341e0e to 436eeb9 Compare September 1, 2026 07:03